import subprocess, m3u8, os, requests, tempfile, sys, uuid, argparse
from Crypto.Cipher import AES
# 读取环境变量 COVER_DIR
COVER_DIR = os.getenv("COVER_DIR", "")
base_dir = os.path.join('./', "public/video-covers")

def m3u8_cover(m3u8_url, out_jpg="cover.jpg"):
    # 1. 拉播放列表
    pls = m3u8.load(m3u8_url)
    seg = pls.segments[0]                       # 只要第一个分片
    ts_url = requests.compat.urljoin(m3u8_url, seg.uri)
    # 2. 下载解密 key
    key_uri = requests.compat.urljoin(m3u8_url, pls.keys[0].uri)
    key = requests.get(key_uri, timeout=10).content
    iv = bytes.fromhex(pls.keys[0].iv[2:]) if pls.keys[0].iv else b'\x00'*16
    # 3. 临时文件：解密后 ts
    fd, tmp_ts = tempfile.mkstemp(suffix='.ts')
    os.close(fd)
    # 4. 边解密边写盘（内存友好）
    cipher = AES.new(key, AES.MODE_CBC, iv)
    with requests.get(ts_url, stream=True, timeout=20) as r, open(tmp_ts, 'wb') as f:
        for chunk in r.iter_content(chunk_size=4096):
            f.write(cipher.decrypt(chunk))
    final_cover_path = os.path.join(base_dir, out_jpg)
    # 5. 截首帧
    subprocess.run([
        "ffmpeg", "-i", tmp_ts,
        "-vframes", "1", "-q:v", "2", "-y", final_cover_path
    ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    os.remove(tmp_ts)
    # print("✅ 封面已保存 ->", os.path.abspath(final_cover_path))
    return f"/video-covers/{out_jpg}"

# if __name__ == "__main__":
    # url = "https://d10cvfvt18g9rw.cloudfront.net/api/app/vid/h5/m3u8/bktappupdata/sp/1p/hr/lz/pm/2b6f5eb547544737a68770b6f6e59e0c.m3u8?token=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ0aW1lc3RhbXAiOjE3NjE3MTc2MDAzMTYxOTMzMDAsInVpZCI6MTgyNTE2MDR9.16w6EyniC2YNKEO4cZ05fqQItQS35a4VO2zfdD424fE&c="
    # 读取命令行参数 分别是url uuid 默认随机生成一个uuid
# parser = argparse.ArgumentParser()
# parser.add_argument('url')
# parser.add_argument('jpg')
# args = parser.parse_args()
# if not args.url or not args.jpg:
#     sys.exit(1)
# res_url = m3u8_cover(args.url, f"video_{args.jpg}_cover.jpg")
# print(res_url)
    # if len(sys.argv) > 1:
    #     url = sys.argv[1]
    #     uuid = sys.argv[2] if len(sys.argv) > 2 else str(uuid.uuid4())
    #     print('生成封面:', url, uuid) 
    #     m3u8_cover(url, f"video_{uuid}_cover.jpg")
